#include <mpi.h>
#include <fstream>
#include <sstream>
#include <assert.h>
#include <math.h>

#include "data.h"
#include "util.h"
#include "gzstream.h"

////////////////////////////////////////////////////////////////////////////////
// Clique
////////////////////////////////////////////////////////////////////////////////
BaseClique::BaseClique()
 : m_first_labels(0), m_second_labels(0), m_gold(-1,-1) 
{}

BaseClique::~BaseClique() 
{}

Clique::Clique() 
{}

Clique::Clique(size_t first_labels, size_t second_labels) 
{
	m_first_labels = first_labels;
	m_second_labels = second_labels;

    m_features = boost::shared_array<SparseArrayPtr>
		(new SparseArrayPtr [m_first_labels*m_second_labels]);
}

Clique::Clique(const Clique &c)
{
    m_first_labels = c.m_first_labels;
    m_second_labels = c.m_second_labels; 
    m_features.reset(new SparseArrayPtr [m_first_labels*m_second_labels]);
    m_gold = c.m_gold;
    for (int x=0; x < m_first_labels*m_second_labels; ++x)
        m_features[x].reset(new SparseArray(*c.m_features[x]));
}

void Clique::read(std::istream &in, int num_features)
{
    std::string buffer;
    int index;
    char tmp_char;
    double value;

    // create the array to store the features
	std::getline(in, buffer, '\n');
	std::istringstream ss(buffer);
    ss >> m_first_labels >> m_second_labels;

	if (m_features.get() == 0)
		m_features = boost::shared_array<SparseArrayPtr>
			(new SparseArrayPtr [m_first_labels*m_second_labels]);

    // read the features
    for(int j=0; j < m_first_labels; ++j)
        for(int i=0; i < m_second_labels; ++i)
        {
            assert(!in.eof());
            std::getline(in, buffer, '\n');
			ss.clear();
            ss.str(buffer);
            SparseArrayPtr sa(new SparseArray(num_features));
            std::vector< std::pair<int,double> > vec;
            while(ss)
            {
                ss >> index;
                if(!ss) break;

                ss.get(tmp_char); assert(tmp_char == ':');
                ss >> value; assert(index < num_features);
                vec.push_back( std::pair<int,double>(index, value));
            }
            std::sort(vec.begin(), vec.end());
            for(std::vector<std::pair<int,double> >::iterator it=vec.begin(); 
                    it != vec.end(); ++it)
                sa->insert(it->first, it->second);

            m_features[j*m_second_labels+i] = sa;
        }
}

void Clique::skip(std::istream &in)
{
    std::string buffer;

    // create the array to store the features
	std::getline(in, buffer, '\n');
	int first_labels, second_labels;
	std::istringstream ss(buffer);
    ss >> first_labels >> second_labels;

    // read the features
    for(int j=0; j < first_labels; ++j)
        for(int i=0; i < second_labels; ++i)
        {
            assert(!in.eof());
            std::getline(in, buffer, '\n');
        }
}

////////////////////////////////////////////////////////////////////////////////
// Node
////////////////////////////////////////////////////////////////////////////////
void Node::read(std::istream &in, int num_features)
{
    std::string buffer;
    int index;
    char tmp_char;
    double value;

    // create the array to store the features
	int labels=0;
    in >> labels;

    // read the features
	for(int i=0; i < labels; ++i)
	{
		assert(!in.eof());

		// detect gold labelling
		bool is_gold=false;
		in >> is_gold;
		if (is_gold) gold=i;

		std::getline(in, buffer, '\n');
		std::istringstream ss(buffer);

		SparseArrayPtr sa(new SparseArray(num_features));
		std::vector< std::pair<int,double> > vec;
		while(ss)
		{
			ss >> index;
			if(!ss) break;

			ss.get(tmp_char); assert(tmp_char == ':');
			ss >> value; assert(index < num_features);
			vec.push_back( std::pair<int,double>(index, value));
		}
		std::sort(vec.begin(), vec.end());
		for(std::vector<std::pair<int,double> >::iterator it=vec.begin(); 
				it != vec.end(); ++it)
			sa->insert(it->first, it->second);

		features.push_back(sa);
//		std::cout << *sa << std::endl;
	}
}

void Node::skip(std::istream &in)
{
    std::string buffer;

	// read labels
	int labels=0;
	std::getline(in, buffer, '\n');
	std::istringstream ss(buffer);
    ss >> labels;

    // read the features
	for(int i=0; i < labels; ++i)
	{
		assert(!in.eof());
		std::getline(in, buffer, '\n');
	}
}

////////////////////////////////////////////////////////////////////////////////
//  Event
////////////////////////////////////////////////////////////////////////////////
Event::Event(size_t o, size_t l)
: order(o), length(l)
{
	// construct the event data structures, initially indexed on segment
	// size (order)
	nodes.resize(order);
	cliques.resize(order);
	for(int o=0; length-o > 0 && o < (int)order; ++o)
	{
		nodes[o].resize(length-o);
		cliques[o].resize((length-1)-o);
	}
}

Event::Event(std::istream &in, int num_features)
{
    in >> order >> length;
//	std::cout << "EVENT: " << order << " " << length << std::endl;
	// construct the event data structures, initially indexed on segment
	// size (order)
	nodes.resize(order);
	cliques.resize(order);
	for(int o=0; length-o > 0 && o < (int)order; ++o)
	{
		nodes[o].resize(length-o);
		cliques[o].resize((length-1)-o);
	}
	// read the node features
    for(int i=0; i < (int)length; ++i)
		for(int o=0; i-o >=0 && o < (int)order; ++o)
		{	
//			std::cout << "\tNODE: " << i << " " << o << std::endl;
			NodePtr n(new Node);
			n->read(in, num_features);
			nodes[o][i-o] = n;	
		}
	// read the clique features
    for(int i=0; i < (int)length-1; ++i)
		for(int o=0; i-o >= 0 && o < (int)order; ++o)
		{	
//			std::cout << "\tCLIQUE: " << i << " " << o << std::endl;
			BaseCliquePtr c(new Clique());
			c->read(in,num_features);
			cliques[o][i-o] = c;
		}
} 

void Event::skip(std::istream &in)
{
	int order=0, length=0;
    std::string buffer;
	std::getline(in, buffer, '\n');
	std::istringstream ss(buffer);
    ss >> order >> length;

	// read the node features
    for(int i=0; i < (int)length; ++i)
		for(int o=0; i-o >=0 && o < (int)order; ++o)
			Node::skip(in);
	// read the clique features
    for(int i=0; i < (int)length-1; ++i)
		for(int o=0; i-o >= 0 && o < (int)order; ++o)
			Clique::skip(in);
} 

////////////////////////////////////////////////////////////////////////////////
//  Event
////////////////////////////////////////////////////////////////////////////////
Events::Events() 
 : m_cliques(0), m_nodes(0)
{ 
    m_init = false;
    m_mpi_start = m_mpi_end = m_curr_event = -1;
}
    
Events::~Events() 
{ 
}

EventPtr Events::event(int c)
{
    assert(c >= m_mpi_start && c < m_mpi_end);
    return m_events[c-m_mpi_start];
}

EventPtr Events::next()
{
    if (m_curr_event < 0)
        m_curr_event = 0;

    if (m_curr_event >= m_mpi_end-m_mpi_start)
    {
        m_curr_event = -1;
        return EventPtr();
    }
    return m_events[m_curr_event++];
}

bool Events::read(const std::string &file, bool disable_mpi)
{
    int num_events=0;

    int rank=0,size=1;
	if (!disable_mpi)
	{
		MPI_Comm_rank(MPI_COMM_WORLD, &rank);
		MPI_Comm_size(MPI_COMM_WORLD, &size);
	}

    // open the file
//    std::ifstream fd(file.c_str());
    igzstream fd(file.c_str());
    if(!fd)
    {
        if (rank == 0) std::cerr << "unable to open input file: " << file<< "\n";
        return false;
    }
     
    // read the header info
    std::string type;
    fd >> type; 
    fd >> num_events >> m_features;

    // set the segment of the data that each MPI process should load
    m_mpi_start = (num_events / size)*rank;
    if (rank == size-1) m_mpi_end = num_events;
    else m_mpi_end = (num_events / size)*(rank+1);
            
    // read each clique entry, one per line
    for(int i=0; i < m_mpi_end; ++i)
    {
        // discard events before the start of our MPI segment
        if (i < m_mpi_start)
		{		
            Event::skip(fd);
		}
        else
        {
            if (rank == 0 && i != 0 && (i % 100) == 0)
            {
                std::cout << '.';
                std::cout.flush();
            }
            if (rank == 0 && i!=0 && (i % 7500) == 0)
                std::cout << " " << i << std::endl;
            EventPtr e(new Event(fd, m_features));
            m_nodes += e->nodes[0].size();
            m_events.push_back(e);
        }
    }
    if (rank==0) std::cout << std::endl;
    fd.close();

    assert(events() == m_mpi_end-m_mpi_start);
//    std::cout << "Rank " << rank << " got " << events() << " events." << std::endl;

    return true;
}

////////////////////////////////////////////////////////////////////////////////
//  FlatEvent
////////////////////////////////////////////////////////////////////////////////
FlatEvents::FlatEvents() 
 : Events()
{
	m_init=false;
}
    
FlatEvents::~FlatEvents() 
{}

bool FlatEvents::read(const std::string &file, bool disable_mpi)
{
    int num_events=0;

    int rank=0,size=1;
	if (!disable_mpi)
	{
		MPI_Comm_rank(MPI_COMM_WORLD, &rank);
		MPI_Comm_size(MPI_COMM_WORLD, &size);
	}

    // open the file
    igzstream fd(file.c_str());
    if(!fd)
    {
        if (rank == 0) std::cerr << "unable to open input file: " << file<< "\n";
        return false;
    }
     
    // read the header info
    std::string type;
    fd >> type; 
    fd >> num_events >> m_features;

    // set the segment of the data that each MPI process should load
    m_mpi_start = (num_events / size)*rank;
    if (rank == size-1) m_mpi_end = num_events;
    else m_mpi_end = (num_events / size)*(rank+1);

	// setup data
	m_num_events = m_mpi_end - m_mpi_start;
	long active_features = process_events(fd, false);
    fd.close();
	m_matrix = SparseMatrixPtr
		(new SparseMatrix(m_num_events*m_num_labels,m_features,active_features)); 
	
	// read the data into the matrix
    igzstream fd2(file.c_str());
    fd2 >> type >> num_events >> m_features;
	process_events(fd2, true);
    fd2.close();

    assert(events() == m_mpi_end-m_mpi_start);
//    std::cout << "Rank " << rank << " got " << events() << " events." << std::endl;
//	std::cout << *m_matrix << std::endl;

    return true;
}

long FlatEvents::process_events(std::istream &fd, bool add_to_matrix)
{
    int rank=0;
	MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    // read each entry
	std::string buffer;
	std::getline(fd, buffer, '\n');
	int order=0, length=0; long num_active=0;
//	std::cout << m_mpi_start << " " << m_mpi_end << std::endl;
    for(int i=0; i < m_mpi_end; ++i)
    {
		// print debug
		if (rank == 0 && (i % 100) == 0)
		{
			std::cout << '.';
			std::cout.flush();
		}
		if (rank == 0 && i!=0 && (i % 7500) == 0)
			std::cout << " " << i << std::endl;

		// read the header info
		fd >> order >> length >> m_num_labels;
		if (order!=1 && length!=1)
		{
			assert(order==1 && length==1);
		}

		// discard the newline
		std::getline(fd, buffer, '\n');
		
		// read the instances
		bool gold; int index; double value; char tmp_char;
		int matrix_index=(i-m_mpi_start)*m_num_labels;
		for (int l=0; l<m_num_labels; ++l)
		{
			std::getline(fd, buffer, '\n');

			// discard events before the start of our MPI segment
			if (i >= m_mpi_start)
			{
				std::istringstream ss(buffer);
				ss >> gold;
				if (add_to_matrix && gold) m_labels.push_back(l);	
				// read the features
				while(ss)
				{
					ss >> index;
					if(!ss) break;
					ss.get(tmp_char); assert(tmp_char == ':');
					ss >> value; assert(index < m_features);
					num_active++;	
					if (add_to_matrix) m_matrix->set(matrix_index+l, index, value);	
				}
			}
		}
    }
	if (add_to_matrix) m_matrix->freeze();	
    if (rank==0) std::cout << std::endl;

	return num_active;;
}
